from mpi4py import MPI
import numpy as np

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

if rank == 0:
    x = np.linspace(0, 100, 12)
else:
    x = None

if rank == 2:
    x_local = np.zeros(9)
else:
    x_local = np.zeros(1)

# if rank == 0:
print("Scatter")
comm.Scatterv([x, (1, 1, 9, 1), (0, 1, 2, 3), MPI.DOUBLE], x_local, root=0)
print(f"Process {rank} has {x_local}")

if rank == 0:
    xGathered = np.zeros(12)
else:
    xGathered = None

print("Gather")
comm.Gatherv(x_local, [xGathered, (1, 1, 9, 1), (0, 1, 2, 3), MPI.DOUBLE], root=0)
print(f"Process {rank} has {xGathered}")
